﻿//-----------------------------------------------------------------------------------------------------------
// WBFSSync Project by Omega Frost 
// http://wbfssync.codeplex.com/
//
// WBFSSync is Licensed under the terms of the 
// Microsoft Reciprocal License (Ms-RL)
//-----------------------------------------------------------------------------------------------------------
using System;
using System.IO;
using System.Runtime.InteropServices;
using WBFSe3.Native;
using WBFSe3.IO;

namespace WBFSe3.Streams
{
    //-------------------------------------------------------------------------------------------------------
    //
    //-------------------------------------------------------------------------------------------------------
    public class WbfsDeviceStream : WbfsStream
    {
        //---------------------------------------------------------------------------------------------------
        // Variables
        //---------------------------------------------------------------------------------------------------
        protected String name;
        protected IntPtr fileHandle;
        protected FileAccess access;

        protected Byte[] bBuffer = null;

        protected bool bNoCache = false;
        protected bool bSynched = false;
        protected long lPosition = 0;
        protected long lSize = 0;
        protected int iOffset = 0;
        protected int iSector = 0;

        private SDISKGEOMETRY dg;
        private SPARTITIONINFO pi;


        //---------------------------------------------------------------------------------------------------
        // Properties
        //---------------------------------------------------------------------------------------------------

        public override string Name { get { return this.name; } }
        public IntPtr Handle { get { return this.fileHandle; } }

        public override bool CanRead { get { return (this.access | FileAccess.Read) != 0; } }
        public override bool CanSeek { get {return true; } }
        public override bool CanWrite { get { return (this.access | FileAccess.Write) != 0; } }

        public override long Length { get { return this.lSize; } }
        public override long Position { get { return this.lPosition; } set { Seek(value, SeekOrigin.Begin); } }

        public SPARTITIONINFO PartitionInfo { get { return this.pi; } }
        public SDISKGEOMETRY DiskGeometry { get { return this.dg; } }


        //---------------------------------------------------------------------------------------------------
        // Routines
        //---------------------------------------------------------------------------------------------------

        
        //---------------------------------------------------------------------------------------------------
        // Constructor
        //---------------------------------------------------------------------------------------------------
        public WbfsDeviceStream(String drive, int sectorSize, FileAccess access)
            : base(sectorSize)
        {
            //Open the drive
            IntPtr ptr = WbfsNative.OpenDrive(drive[0], access);

            //Get the size
            SDISKGEOMETRY dg = WbfsNative.GetDiskGeometry(ptr);
            SPARTITIONINFO pi = WbfsNative.GetPartitionInfo(ptr);

            this.name = Char.ToUpper(drive[0]) + ":\\";
            this.fileHandle = ptr;
            this.access = access;
            this.lSize = pi.PartitionLength;
            this.sectorSize = sectorSize;
            this.bBuffer = new Byte[sectorSize];
            this.bSynched = true;
            this.dg = dg;
            this.pi = pi;

            Buffer();
        }


        //---------------------------------------------------------------------------------------------------
        // Constructor
        //---------------------------------------------------------------------------------------------------
        public WbfsDeviceStream(String drive, FileAccess access)
            : base(0)
        {
            //Open the drive
            IntPtr ptr = WbfsNative.OpenDrive(drive[0], access);

            //Get the size
            SDISKGEOMETRY dg = WbfsNative.GetDiskGeometry(ptr);
            SPARTITIONINFO pi = WbfsNative.GetPartitionInfo(ptr);

            this.name = Char.ToUpper(drive[0]) + ":\\";
            this.fileHandle = ptr;
            this.access = access;
            this.lSize = pi.PartitionLength;
            this.sectorSize = dg.BytesPerSector;
            this.bBuffer = new Byte[dg.BytesPerSector];
            this.bSynched = true;
            this.dg = dg;
            this.pi = pi;

            Buffer();
        }


        //---------------------------------------------------------------------------------------------------
        // Seeks to a new position of the disk
        //---------------------------------------------------------------------------------------------------
        public override long Seek(long offset, SeekOrigin origin)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            switch (origin)
            {
                case SeekOrigin.Current:
                    offset += this.lPosition;
                    break;
                case SeekOrigin.End:
                    offset = this.lSize - offset;
                    break;
            }

            if (offset < 0)
                throw new ArgumentOutOfRangeException();

            if (offset > this.lSize)
                throw new EndOfStreamException();

            int newsector = (int)(offset / this.sectorSize);
            this.iOffset = (int)(offset % this.sectorSize);
            this.lPosition = offset;

            if (newsector != this.iSector)
            {
                this.iSector = newsector;

                if (!this.bNoCache)
                    Sync();

                //Seek to the begin of the sector

                long dummypos;
                if (!WbfsNative.SetFilePointerEx(this.fileHandle, 
                    (long)newsector * this.sectorSize,
                    out dummypos, SeekOrigin.Begin))
                    throw new WbfsIOException();

                if (!this.bNoCache)
                    Buffer();
            }

            return offset;
        }


        //---------------------------------------------------------------------------------------------------
        // Do nothing...
        //---------------------------------------------------------------------------------------------------
        public override void SetLength(long value)
        {
            //Nops...
        }


        //---------------------------------------------------------------------------------------------------
        // Reads an array from the buffer/disk
        //---------------------------------------------------------------------------------------------------
        public override int Read(byte[] array, int offset, int count)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (array.Length < offset + count)
                throw new ArgumentException();

            if (array == null)
                throw new ArgumentNullException();

            if ((offset < 0) || (count < 0))
                throw new ArgumentOutOfRangeException();

            if (!this.CanRead)
                throw new NotSupportedException();

            //Prevents from reading beyond EOF
            if (this.lPosition >= this.lSize)
                throw new EndOfStreamException();

            int icount = 0;
            while (count > 0)
            {
                //Calculate bytes to read
                int max = this.sectorSize - this.iOffset;
                int toread = count < max ? count : max;

                //Copy the data between buffers
                Array.Copy(this.bBuffer, this.iOffset, array, 
                    offset + icount, toread);

                //Recalculate position
                icount += toread;
                count -= toread;
                this.lPosition += toread;
                this.iOffset += toread;
                
                //Sync sector, Buffer next
                if (this.iOffset >= this.sectorSize)
                {
                    Sync();
                    
                    this.iSector++;
                    this.iOffset = 0;

                    if (this.lPosition == this.lSize)
                        break;

                    Buffer();
                }
            }

            return icount;
        }


        //---------------------------------------------------------------------------------------------------
        // Writes an array to the buffer/disk
        //---------------------------------------------------------------------------------------------------
        public override void Write(byte[] array, int offset, int count)
        {
            if (this.fileHandle == WbfsNative.InvalidHandle)
                throw new ObjectDisposedException("WbfsFileStream");

            if (array.Length < offset + count)
                throw new ArgumentException();

            if (array == null)
                throw new ArgumentNullException();

            if ((offset < 0) || (count < 0))
                throw new ArgumentOutOfRangeException();

            if (!this.CanWrite)
                throw new NotSupportedException();

            //Prevents from reading beyond EOF
            if (this.lPosition >= this.lSize)
                throw new EndOfStreamException();

            int icount = 0;
            while (count > 0)
            {
                //Mark as modified
                this.bSynched = false;

                //Calculate bytes to read
                int max = this.sectorSize - this.iOffset;
                int towrite = count < max ? count : max;

                //Copy the data between buffers
                Array.Copy(array, offset + icount, 
                    this.bBuffer, this.iOffset, towrite);
                
                //Recalculate position
                icount += towrite;
                count -= towrite;
                this.lPosition += towrite;
                this.iOffset += towrite;

                //Sync sector, Buffer next
                if (this.iOffset >= this.sectorSize)
                {
                    Sync();

                    this.iSector++;
                    this.iOffset = 0;

                    if (this.lPosition == this.lSize) break;

                    Buffer();
                }
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Reads an aligned array to the disk
        //---------------------------------------------------------------------------------------------------
        public override int ReadSector(int sector, IntPtr array, int offset)
        {
            uint read;

            if (!WbfsNative.SetFilePointerEx(this.fileHandle,
                    (long)this.sectorSize * sector,
                    out this.lPosition, SeekOrigin.Begin))
                throw new WbfsIOException();

            this.iSector = sector;

            if (!WbfsNative.ReadFile(this.fileHandle, array,
                (uint)this.sectorSize, out read, IntPtr.Zero))
                throw new WbfsIOException();

            this.lPosition += read;
            return (int)read;
        }


        //---------------------------------------------------------------------------------------------------
        // Writes an aligned array to the disk
        //---------------------------------------------------------------------------------------------------
        public override int WriteSector(int sector, IntPtr array, int offset)
        {
            uint written;

            if (!WbfsNative.SetFilePointerEx(this.fileHandle,
                    (long)this.sectorSize * sector,
                    out this.lPosition, SeekOrigin.Begin))
                throw new WbfsIOException();

            this.iSector = sector;

            if (!WbfsNative.WriteFile(this.fileHandle, array,
                (uint)this.sectorSize, out written, IntPtr.Zero))
                throw new WbfsIOException();

            this.lPosition += written;
            return (int)written;
        }


        //---------------------------------------------------------------------------------------------------
        // Writes empty data to the disk
        //---------------------------------------------------------------------------------------------------
        public override void WriteEmptyBlock(int count)
        {
            if (this.lPosition + count == this.lSize)
                throw new EndOfStreamException();

            Byte[] zeroes = new Byte[count];
            Array.Clear(zeroes, 0, zeroes.Length);

            if (this.bNoCache) Write(zeroes, 0, count);
            else
            {
                uint written;

                if (!WbfsNative.WriteFile(this.fileHandle, zeroes,
                    (uint)this.sectorSize, out written, IntPtr.Zero))
                    throw new WbfsIOException();

                this.lPosition += written;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Updates the disk with the buffer data
        //---------------------------------------------------------------------------------------------------
        public override void Flush()
        {
            Sync();
        }


        //---------------------------------------------------------------------------------------------------
        // Closes the stream
        //---------------------------------------------------------------------------------------------------
        public override void Close()
        {
            base.Close();
        }


        //---------------------------------------------------------------------------------------------------
        // Closes the handle
        //---------------------------------------------------------------------------------------------------
        protected override void Dispose(bool disposing)
        {
            WbfsNative.CloseHandle(this.fileHandle);
            this.fileHandle = WbfsNative.InvalidHandle;
            base.Dispose(disposing);
        }


        //---------------------------------------------------------------------------------------------------
        // Synchronizes the data with the disk
        //---------------------------------------------------------------------------------------------------
        private void Sync()
        {
            if (this.bSynched)
                return;

            //Seek to the begin of the sector

            long dummypos;
            if (!WbfsNative.SetFilePointerEx(this.fileHandle, -this.sectorSize,
                out dummypos, SeekOrigin.Current))
                throw new WbfsIOException();

            //Write data

            uint write;
            if (!WbfsNative.WriteFile(this.fileHandle, this.bBuffer,(uint)this.sectorSize,
                out write, IntPtr.Zero) || (write != this.sectorSize))
                throw new WbfsIOException();

            this.bSynched = true;
        }


        //---------------------------------------------------------------------------------------------------
        // Buffer a sector from the disk
        //---------------------------------------------------------------------------------------------------
        private void Buffer()
        {
            uint read;
            if (!WbfsNative.ReadFile(this.fileHandle, this.bBuffer, (uint)this.sectorSize,
                out read, IntPtr.Zero) || (read != this.sectorSize))
                throw new WbfsIOException();
        }
    }
}
